import os
import torch
# import torch.nn as nn
# from torchsummary import summary
from torchvision.models import mobilenet_v3_small, MobileNet_V3_Small_Weights


def mobilenet_v3_s_model(classes_num=5, download=False, freeze=False, mode='train'):
    """
    :param classes_num: 数据的类别个数
    :param download: 是否从官网上下载预训练权重
    :param freeze: 是否冻结权重
    :param mode: {'train': 模型训练使用, 'predict': 模型预测使用}
    :return: efficient_v2_s_model
    """
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    print('device: ', device)

    if mode == 'train':
        if download:  # 若`download=True`, 则会自动将权重下载导本地用户缓存(会显示路径), 下载成功与否取决于当前网络
            # 第一次加载会自动下载权重
            # 查看该权重要求的输入, 主要查看的有：[crop_size, mean, std], 读取自己的数据时要用到
            print('train transforms requires: \n', MobileNet_V3_Small_Weights.IMAGENET1K_V1.transforms())
            model = mobilenet_v3_small(weights=MobileNet_V3_Small_Weights.IMAGENET1K_V1)
        else:
            # 权重所在路径
            weights_path = './weights/mobilenet_v3_small_pretrain.pth'
            assert os.path.exists(weights_path), \
                '请将预训练权重放至当前工程的`weights`文件夹下, 并将文件名改为`mobilenet_v3_small_pretrain.pth`'
            weights = torch.load(weights_path)  # 读取权重, 这里加载到
            model = mobilenet_v3_small()         # 实例化模型
            model.load_state_dict(state_dict=weights)  # 模型加载权重
    if mode == 'predict':
        model = mobilenet_v3_small()  # 实例化模型

    # 修改模型的输出层(主要修改输出类别数量)
    in_channel = model.classifier[3].in_features
    # 直接将输出类别数改成我们自己数据集的类别数, 我们`classes_num=5`
    model.classifier[3] = torch.nn.Linear(in_features=in_channel, out_features=classes_num)

    if freeze:
        # 冻结权重, 仅训练最后的`fc`层, 但模型效果较差; 若训练全部层, 效果较好, 不过电脑内存要求比较高
        for name, param in model.named_parameters():
            if 'classifier' not in name:
                param.requires_grad = False

    return model.to(device)

